import copy
import math

import common_modules

from DrawingAPIs.breakdown_drawer import *

# Read the memory overhead occupied by column buffers required by the convolution operation.
# File: <device>/IntraLayer/<network>/BuffOverhead/ColBuff-MEM.csv
def ColBuffMemData(col_buff_file):
    col_buff_mem = 0.0
    with open(col_buff_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            col_buff_mem += float(row[3][:-3])
    return col_buff_mem

# Read the CUPTI buffer overhead allocated by the CUPTI runtime.
# File: <device>/IntraLayer/<network>/CuptiOverhead/CUPTI-BUFFER.csv
def CuptiBuffMemData(cupti_buff_file, unit='MB'):
    cupti_buff_mem = 0.0
    with open(cupti_buff_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            cupti_buff_mem += int(row[0][:-3])
    if unit == 'MB':
        cupti_buff_mem /= 1024.0

    return cupti_buff_mem

# Read temp buffer overhead for kernels' timestamp and configuration
# File: <device>/IntraLayer/<network>/CuptiOverhead/TEMP-BUFFER-OVERHEAD.csv
def TempBuffOverheadData(temp_buff_file, unit='MB'):
    timestamp_buff_mem = 0.0
    kernel_buff_mem = 0.0
    with open(temp_buff_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            timestamp_buff_mem += (float(row[1]) / 1024.0)
            kernel_buff_mem += (float(row[3]) / 1024.0)

    if unit == 'MB':
        timestamp_buff_mem /= 1024.0
        kernel_buff_mem /= 1024.0

    return timestamp_buff_mem, kernel_buff_mem

def DrawBuffOverheadFig(args, devices):
    col_buff_overhead = {}
    cupti_buff_overhead = {}
    timestamp_buff_overhead = {}
    kernel_buff_overhead = {}
    for dev in devices:
        col_buff_overhead.setdefault(dev, [])
        cupti_buff_overhead.setdefault(dev, [])
        timestamp_buff_overhead.setdefault(dev, [])
        kernel_buff_overhead.setdefault(dev, [])

        curr_root_folder = args.root[0] + '/' + dev + '/IntraLayer'
        nets = ['CIFAR10_nv', 'AlexNet', 'Resnet18', 'VGG16']
        y_max_val = 0.0
        # for net in os.listdir(curr_root_folder):
        for net in nets:
            # if net.lower() == "googlenet":
            #     continue
            # if net not in nets:
            #     nets.append(net)
            col_buff_file = curr_root_folder + '/' + net + '/BuffOverhead/ColBuff-MEM.csv'
            cupti_buff_file = curr_root_folder + '/' + net + '/CuptiOverhead/CUPTI-BUFFER.csv'
            temp_overhead_file = curr_root_folder + '/' + net + '/CuptiOverhead/TEMP-BUFFER-OVERHEAD.csv'

            # col_buff_mem = ColBuffMemData(col_buff_file)
            cupti_buff_mem = CuptiBuffMemData(cupti_buff_file)
            timestamp_buff_mem, kernel_buff_mem = TempBuffOverheadData(temp_overhead_file)

            # col_buff_overhead[dev].append(col_buff_mem)
            cupti_buff_overhead[dev].append(cupti_buff_mem)
            timestamp_buff_overhead[dev].append(timestamp_buff_mem)
            kernel_buff_overhead[dev].append(kernel_buff_mem)

            if y_max_val < (cupti_buff_mem + timestamp_buff_mem + kernel_buff_mem):
            # if y_max_val < (col_buff_mem + cupti_buff_mem + timestamp_buff_mem + kernel_buff_mem):
                y_max_val = cupti_buff_mem + timestamp_buff_mem + kernel_buff_mem;
                # y_max_val = col_buff_mem + cupti_buff_mem + timestamp_buff_mem + kernel_buff_mem;

        y_value_set = [cupti_buff_overhead[dev], timestamp_buff_overhead[dev],
                       kernel_buff_overhead[dev]]
        # y_value_set = [col_buff_overhead[dev], cupti_buff_overhead[dev],
        #                timestamp_buff_overhead[dev], kernel_buff_overhead[dev]]
        # legend_labels = ["$mem_{aux}$", "$mem_{cupti}$", "$mem_{tt}$", "$mem_{K}$"]
        legend_labels = ["$mem_{cupti}$", "$mem_{tt}$", "$mem_{K}$"]
        if y_max_val > 999:
            y_max_val = math.ceil(y_max_val / 100.0) * 100
        elif y_max_val > 100 and y_max_val < 1000:
            y_max_val = math.ceil(y_max_val / 50.0) * 50
        elif y_max_val < (math.ceil(y_max_val) - 0.5):
            y_max_val = math.ceil(y_max_val) - 0.5
        else:
            y_max_val = math.ceil(y_max_val)
        # print("output folder: {0}".format(args.output[0] + '/MemOverhead'))
        DrawBreakdownFig(nets, y_value_set, legend_labels, "Neural Network Model", "Memory (MB)",
                         x_min=0.5, x_max=(len(nets) + 0.5), y_min=0, y_max=y_max_val, bar_width=0.4,
                         folder=args.output[0] + '/MemOverhead', filename=dev + '_exp_buffCost_total_mem',
                         allow_legend=True)
